# -*- coding: utf-8 -*-
"""
Created on Fri Dec 10 15:08:55 2021

@author: zhiyinan
"""

import numpy as np
import mpctools as mpc
import matplotlib.pyplot as plt
import tensorflow as tf
import Parameters1D_casadi as par 
from Model1D_casadi_crop import *
import van_Genuchten1D_casadi as vg 
from scipy import integrate
from sklearn.preprocessing import MinMaxScaler
from casadi import *
import matplotlib.font_manager as font_manager
import time
import pandas as pd
from sklearn.linear_model import Ridge
import random


totalDepth, axialNodes, axialDistance, totalNodes = par.spatialVariables_1D()  #Spatial parameters
samplingTime, samplingTimeInternal, internalTimeSteps, totTimeSteps = \
                                 par.temporalVariable_dataGen(25000*6*60)

soilPars = loamySoil()

Nu = 1
Nx = 26
Ny = 1
Ns = 1
Np = 2


Delta = 2*60*60
Nt = 40*60*60
Nsim = 5*24*60*60
def ODE(head, irrigAmount, p):
    cropCoeff = p[0]
    refEvap = p[1]
    origin = Richards1D_casadi(head, irrigAmount, cropCoeff, refEvap)
    error = DM.zeros(26)
    # for i in range(26):
    #     error[i] = np.random.choice([-1,1]) * 0.1*np.random.random_sample()*origin[i]
    return origin + error
# lfunc, [Ny, Nu, Ns], ["x","u","s"], funcname="l"
ode_sim = mpc.DiscreteSimulator(ODE, Delta, [Nx, Nu, Np], ["x","u","p"])
f = mpc.getCasadiFunc(ODE, [Nx, Nu, Np], ["x","u","p"], funcname="f")

y_index = 19

lbZone = 0.18
ubZone = 0.23

Q = 1000
R = 5
lbu = -2e-05  #  unit [m/s]
ubu = 0


h0 = -0.96
x0 = np.ones(26)*h0
y0 = vg.volumetricMoisture(x0[y_index])

cropCoeff = 0.88*DM.ones(int(Nt/Delta))
refEvap = 3.102e-08*DM.ones(int(Nt/Delta))
para = horzcat(cropCoeff, refEvap) 

# ========================Define Objective function=======================

# def obj(u, x, yz, Q, R):
#     y = vg.volumetricMoisture(x[y_index])
#     J = mtimes((y - yz).T, mtimes(Q, (y - yz))) + R*u**2
#     return J

# def Pffunc(x):
#     dx = y - s
#     return mpc.mtimes(dx.T, Q, dx) + mpc.mtimes(u.T,R,u)
# Pf = mpc.getCasadiFunc(lfunc, [Ny, Nu, Ns], ["x","u","sf"], funcname="Pf")


def lfunc(x,u,s):
    y = vg.volumetricMoisture(x[y_index])
    dx = y - s
    return mpc.mtimes(dx.T, Q, dx) + mpc.mtimes(u.T,R,u)
l = mpc.getCasadiFunc(lfunc, [Nx, Nu, Ns], ["x","u","s"], funcname="l")

def zone_con(s):
    return vertcat(s - ubZone, lbZone - s) 

e = mpc.getCasadiFunc(zone_con, [Ns], ["s"], funcname="e")
# ====================Define MPC controller====================
funcargs = {"f" : ["x","u","p"], "e" : ["s"], "l" : ["x","u","s"]}

contargs = dict(
        N = {"t":int(Nt/Delta), "x":Nx, "u":Nu, "p":Np, "s":Ns, "e":Ns*2, "c":2},
        funcargs = funcargs,
        verbosity=5,
        l=l,
        e=e,
        x0=x0,
        p = para,
        #Pf=Pf,
        ub = {"u":ubu,"x":-0.1*np.ones(Nx)}, # Change upper bounds 
        lb = {"u":lbu,"x":-80*np.ones(Nx)},	# Change lower bounds 
        guess = {
                "x":-0.96*np.ones(Nx),
                "u":np.zeros(Nu)
                },
        )

# Create MPC controller
mpc_controller = mpc.nmpc(f=f,Delta=Delta,**contargs)


x = np.zeros((int(Nsim/Delta)+1, Nx))
y = np.zeros((int(Nsim/Delta)+1, Ny))
yz = np.zeros((int(Nsim/Delta)+1, Ns))
x[0,:] = x0
y[0,:] = y0


u = np.zeros((int(Nsim/Delta), Nu))

#=================================================================#
# Simulate closed loop system
#=================================================================#

for t in range(int(Nsim/Delta)):
    
    # Update intial condition and solve control problem.
    mpc_controller.fixvar("x",0,x[t,:]) # Set initial value
    mpc_controller.solve()
    
    # Print status and make sure solver didn't fail.
#    print "%5s %d: %s" % ("nmpc",t,mpc_controller.stats["status"])
#    if mpc_controller.stats["status"] != "Solve_Succeeded":
#        break
#    else:
#        mpc_controller.saveguess()
        
    u[t,:] = np.squeeze(mpc_controller.var["u",0])
    yz[t,:] = np.squeeze(mpc_controller.var["s",0])
    x[t+1,:] = ode_sim.sim(x[t,:], u[t,:], vertcat(0.88, 3.102e-08))
    y[t+1,:] = vg.volumetricMoisture(x[t+1,y_index])

